# -*- coding:utf-8 -*-  
import os
import datetime
import torch
import torchvision.transforms as transforms
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader
from train import train
import config as c
#os.environ["CUDA_VISIBLE_DEVICES"] = "3"
def main():

    device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')

    #重新训练的文件夹名
    train_dir = "/home/lzc/zzc/HiDDeN/runs-works/CNN/run_Brightness-no-normalize_2024-03-20_01-58-18"

    train_data_dir = "/home/lzc/DRHiNet/data/train"  # 训练数据
    test_data_dir = "/home/lzc/DRHiNet/data/val"     # 测试数据

    model_dir = "model"
    model_dir = os.path.join(train_dir, model_dir)
    #runs-works/name/model
    os.makedirs(model_dir, exist_ok=True)
    
    log_dir = "log"
    log_dir = os.path.join(train_dir, log_dir)
    
    os.makedirs(log_dir, exist_ok=True)

    log_file = os.path.join(log_dir, "log.csv")

    save_image_dir = "image"
    save_image_dir = os.path.join(train_dir, save_image_dir)
    os.makedirs(save_image_dir, exist_ok=True)

    # Super parameter
    image_size = 128
    batch_size = 64
    message_length = 30
    num_epochs = 600

    # 是否进行归一化
    normalize = False

    # 定义图像预处理的转换操作
    data_transforms = {
        'train': transforms.Compose([
            # 针对训练集的数据预处理：
            # 随机裁剪图像到指定尺寸（image_size）并进行填充（pad_if_needed=True）
            transforms.RandomCrop((image_size, image_size), pad_if_needed=True),
            # 将图像转换为张量（Tensor）格式
            transforms.ToTensor(),
            # 如果需要进行图像归一化，将像素值从 [0, 255] 映射到 [-1, 1]
            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]) if normalize else transforms.Compose([]),
        ]),
        'test': transforms.Compose([
            # 针对测试集的数据预处理：
            # 居中裁剪图像到指定尺寸（image_size）
            transforms.CenterCrop((image_size, image_size)),
            # 将图像转换为张量（Tensor）格式
            transforms.ToTensor(),
            # 如果需要进行图像归一化，将像素值从 [0, 255] 映射到 [-1, 1]
            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]) if normalize else transforms.Compose([]),
        ])
    }

    # 创建ImageFolder数据集实例
    train_dataset = ImageFolder(train_data_dir, transform=data_transforms['train'])
    test_dataset = ImageFolder(test_data_dir, transform=data_transforms['test'])

    # 创建数据加载器
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=4)

    en = 0.7
    de = 1.0
    adv = 0.001

    train(device, train_loader, test_loader, message_length, num_epochs, model_dir, log_file, save_image_dir, en, de, adv, normalize)

if __name__ == '__main__':
    main()